import xarray as xr
import math
import numpy as np
import pandas as pd
import sharppy.sharptab.thermo as thermo
import sharppy.sharptab.utils as utils

N=10 
fnewlatlon1="locnontor.txt"
flatlon1=open(fnewlatlon1,"a")
f=open("NToS.txt","r",encoding='utf-8').readlines()
for number in range(143):
    print("*************************************************")
    print(number)
    fa=f[number]
    k=number+1
    year=fa[0:4]
    month=fa[4:6]
    day=fa[6:8]
    time=fa[8:10]
    site=fa[-5:-1]
    late=float(fa[11:16])
    lone=float(fa[17:-6])

    ds=xr.open_dataset('data/NToS_era/ERA5'+fa[0:10]+'.grib',engine='cfgrib')
    dscape=xr.open_dataset('data/NToS_CAPE/ERA5'+fa[0:10]+'.grib',engine='cfgrib')

    lats=math.floor(late*100/25)*25/100-(N/2-1)*0.25   #lat_start
    latd=math.floor(late*100/25)*25/100+N/2*0.25       #lat_end
    lons=math.floor(lone*100/25)*25/100-(N/2-1)*0.25   #lon_start
    lond=math.floor(lone*100/25)*25/100+N/2*0.25       #lon_end

    lat_area,lon_area=np.mgrid[latd:lats:10j,lons:lond:10j]
    dis_data=np.sqrt(np.power(lat_area-late,2)+np.power(lon_area-lone,2))*100
    lat1=np.arange(latd,lats-0.1,-0.25)
    lon1=np.arange(lons,lond+0.1,0.25)
    distance=pd.DataFrame(dis_data,index=lat1,columns=lon1)

    dis=100
    TF=distance.copy()
    TF[distance<=dis]=1
    TF[distance>dis]=0

    u=ds['u']
    v=ds['v']
    uN=(u.loc[1000,latd:lats,lons:lond]+u.loc[975,latd:lats,lons:lond]+u.loc[950,latd:lats,lons:lond])/3
    vN=(v.loc[1000,latd:lats,lons:lond]+v.loc[975,latd:lats,lons:lond]+v.loc[950,latd:lats,lons:lond])/3

    wdirN=np.arctan2(vN, uN)* 180 / np.pi
    for i in range(N):
        for j in range(N):
            if (vN[i,j]==0)&(uN[i,j]==0):
                wdirN[i,j]=-9999

    theta1=np.arctan2(late-lat_area,lone-lon_area)*180/np.pi
    theta=pd.DataFrame(theta1,index=lat1,columns=lon1)

    for i in range(N):
        for j in range(N):
            if (-105<theta.iloc[i,j]) and (theta.iloc[i,j]<=105):
                low=theta.iloc[i,j]-75
                high=theta.iloc[i,j]+75
                if wdirN[i,j]<=high and wdirN[i,j]>=low:
                    TF.loc[float(wdirN[i,j].latitude.values),float(wdirN[i,j].longitude.values)]+=1
            elif theta.iloc[i,j]<=-105:
                low=180-(75-theta.iloc[i,j]-180)
                high=theta.iloc[i,j]+75
                if (wdirN[i,j]<=high and wdirN[i,j]>-180) or (wdirN[i,j]>=low and wdirN[i,j]<=180):
                    TF.loc[float(wdirN[i,j].latitude.values),float(wdirN[i,j].longitude.values)]+=1
            elif theta.iloc[i,j]>105:
                low=theta.iloc[i,j]-75
                high=-180+theta.iloc[i,j]+75-180 
                if (wdirN[i,j]>=low and wdirN[i,j]<=180) or (wdirN[i,j]<=high and wdirN[i,j]>-180):
                    TF.loc[float(wdirN[i,j].latitude.values),float(wdirN[i,j].longitude.values)]+=1
    a=(TF<=1)
    if a.all(axis = None):
        print('continiue')
        continue

    points_lat=[]
    points_lon=[]
    ind=np.where(TF==2)
    cape_all=dscape.cape.loc[latd:lats,lons:lond]
    for op in range(len(ind[0])):  
        points_lat.append(cape_all[ind[0][op],ind[1][op]].latitude.values)
        points_lon.append(cape_all[ind[0][op],ind[1][op]].longitude.values)

    cape_all=dscape.cape.loc[latd:lats,lons:lond]
    sig=0 
    for op in range(len(ind[0])):  
        i = ind[0][op]
        j = ind[1][op]
        if sig>=1:
            if cape_all[i,j]>cape1:
               cape1=cape_all[i,j]
               lat=cape_all[i,j].latitude.values
               lon=cape_all[i,j].longitude.values
        else:
            cape1=cape_all[i,j]
            lat=cape_all[i,j].latitude.values
            lon=cape_all[i,j].longitude.values
        sig=sig+1

    lat5=[lat,lat-0.25,lat+0.25,lat,lat]
    lon5=[lon,lon,lon,lon-0.25,lon+0.25]
    
    lat_5ave=[]
    lon_5ave=[]
    for i in range(len(points_lat)):
        for j in range(len(lat5)):
            if (lat5[j]==points_lat[i] and lon5[j]==points_lon[i]):   
                lat_5ave.append(lat5[j])
                lon_5ave.append(lon5[j])
    line5=str(k)+"  "+fa[0:10]+"  "+str(lat)+","+str(lon) 
    print(lat_5ave)
    if len(lat_5ave)>1:
        for i in range(1,len(lat_5ave)):
            print(i)
            line5=line5+"  "+str(lat_5ave[i])+","+str(lon_5ave[i]) 
    line5=line5+"\n"    
    flatlon1.write(line5)   

    n5=len(lat_5ave)
    u5=u.loc[:,lat_5ave[0],lon_5ave[0]] 
    v5=v.loc[:,lat_5ave[0],lon_5ave[0]]
    T5=ds.t.loc[:,lat_5ave[0],lon_5ave[0]]
    hght5=ds.z.loc[:,lat_5ave[0],lon_5ave[0]] 
    r5=ds.r.loc[:,lat_5ave[0],lon_5ave[0]]
    q5 = ds.q.loc[:, lat_5ave[0], lon_5ave[0]]

    if n5>1:
        for i in range(1,n5):
            u5=u5+u.loc[:,lat_5ave[i],lon_5ave[i]] 
            v5=v5+v.loc[:,lat_5ave[i],lon_5ave[i]] 
            T5=T5+ds.t.loc[:,lat_5ave[i],lon_5ave[i]] 
            hght5=hght5+ds.z.loc[:,lat_5ave[i],lon_5ave[i]] 
            r5=r5+ds.r.loc[:,lat_5ave[i],lon_5ave[i]]             
    u5=u5/n5 
    v5=v5/n5
    wdir5,wspd5=utils.comp2vec(u5,v5)
    wspd5=wspd5/0.5144
    T5=T5/n5
    T5=T5-273.15 
    hght5=hght5/n5/9.80665
    r5=r5/n5
    vappres_td=r5/100*thermo.vappres(T5)
    dwpt5=thermo.temp_at_vappres(vappres_td)
        
    titlename=site+"     "+year+month+day+"/"+time+'00'+"  "+str(late)+','+str(lone)
    fnew="NToS_5/"+str(k)+'_'+year+month+day+time+".txt"
    f1=open(fnew,"a")
    f1.write("%TITLE%\n")
    f1.write(titlename+'\n')
    f1.write("LEVEL       HGHT       TEMP       DWPT       WDIR       WSPD\n")
    f1.write("-------------------------------------------------------------------\n")
    f1.write("%RAW%\n")
    for i in range(37):
        f1.write(str(u.isobaricInhPa[i].values)+","+str(hght5[i].values)+","+str(T5[i].values)+","+str(dwpt5[i].values)+","+str(wdir5[i])+","+str(wspd5[i])+"\n")
    f1.write("%END%")
    f1.close()